import torch
from torch import nn
import math

class Embedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.lut = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

emb = Embedding(10,8)
inputs = torch.tensor(
    [
        [1,2,3],
        [4,5,6],
    ]
)
emb_out = emb(inputs)
print(emb_out.shape)
print(emb_out)